package com.fly.consultant.config;

import com.fly.consultant.service.ConsultantService;
import com.fly.consultant.store.RedisChatMemoryStore;
import dev.langchain4j.community.store.embedding.redis.RedisEmbeddingStore;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.loader.ClassPathDocumentLoader;
import dev.langchain4j.data.document.parser.apache.pdfbox.ApachePdfBoxDocumentParser;
import dev.langchain4j.data.document.splitter.DocumentSplitters;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.memory.chat.ChatMemoryProvider;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.memory.chat.TokenWindowChatMemory;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.util.List;

@Configuration
public class CommonConfig {

    @Autowired
    private RedisChatMemoryStore redisChatMemoryStore;
    @Autowired
    private EmbeddingModel embeddingModel;
    @Autowired
    private RedisEmbeddingStore redisEmbeddingStore;


//    // 注入模型
//    @Autowired
//    private OpenAiChatModel openAiChatModel;
//
//
//    // 注入服务
//    @Bean
//    public ConsultantService consultantService(){
//        return AiServices.builder(ConsultantService.class)
//                // 选择模型
//                .chatModel(openAiChatModel)
//                .build();
//    }

    // 注入记忆模型
    @Bean
    public ChatMemory chatMemory(){
        return MessageWindowChatMemory.builder()
                .maxMessages(20)
                .build();
//        return TokenWindowChatMemory.builder().build();
    }

    // 注入会话记忆提供者
    @Bean
    public ChatMemoryProvider chatMemoryProvider(){
        ChatMemoryProvider chatMemoryProvider =  new ChatMemoryProvider() {
            @Override
            public ChatMemory get(Object memoryId) {
                return MessageWindowChatMemory.builder()
                        .id(memoryId)
                        .chatMemoryStore(redisChatMemoryStore)
                        .maxMessages(20)
                        .build();
            }
        };
        return chatMemoryProvider;
    }

    // 创建内存向量数据库对象
//    @Bean
    public EmbeddingStore store(){
        // 导入数据
//        List<Document> content = ClassPathDocumentLoader.loadDocuments("content");
        // 添加对pdf的解析
        List<Document> content = ClassPathDocumentLoader.loadDocuments("content",new ApachePdfBoxDocumentParser());
        // 创建内存向量数据库存储对象
//        InMemoryEmbeddingStore store = new InMemoryEmbeddingStore();
        // 自定义文档分割器
        DocumentSplitter ds = DocumentSplitters.recursive(500,100);
        // 分片存储
        EmbeddingStoreIngestor build = EmbeddingStoreIngestor.builder()
                .embeddingStore(redisEmbeddingStore)
                .embeddingModel(embeddingModel)
                .documentSplitter(ds)
                .build();
        build.ingest(content);
        // 返回
        return redisEmbeddingStore;
    }

    // 创建检索向量数据库对象
    @Bean
    public ContentRetriever contentRetriever(EmbeddingStore store){
        return EmbeddingStoreContentRetriever.builder()
                // 调用的向量数据库对象
                .embeddingStore(store)
                // 检索数量
                .maxResults(3)
                // 检索分数
                .minScore(0.5)
                // 调用向量模型
                .embeddingModel(embeddingModel)
                .build();
    }

}
